# RAG检索

import os
import openai
from IPython.utils import docs
from llama_index import SimpleDirectoryReader, Document, VectorStoreIndex, ServiceContext, load_index_from_storage, \
    StorageContext
from llama_index.llms import OpenAI
from langchain.embeddings.openai import OpenAIEmbeddings
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.node_parser import SentenceWindowNodeParser
import os
from llama_index.postprocessor import MetadataReplacementPostProcessor, SentenceTransformerRerank

api_key = "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"
os.environ['OPENAI_API_KEY'] = "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"

# 1、加载文件
documents = SimpleDirectoryReader(
    input_files=["./合同履约-合同结算管理操作手册.pdf"]
).load_data()

# 2、输入数据（查看是否正常导入）
# print(type(documents))
# print(len(documents))
# print(type(documents[0]))
# print(documents[0])

# 3、合并为单个文档
document = Document(text="\n\n".join([doc.text for doc in documents]))

# 4、对文档进行索引
llm = OpenAI(api_key=api_key, model_name="gpt-3.5-turbo", temperature=0)
service_context = ServiceContext.from_defaults(
    llm=llm, embed_model=OpenAIEmbedding()
)
index = VectorStoreIndex.from_documents([document], service_context=service_context)

# 5、获取查询引擎，允许我们发送用户查询，进行数据检索和综合
query_engine = index.as_query_engine()

# 6、查询
response = query_engine.query("结合PDF中的内容，用中文回复，发票查询应该怎么做？")
print(response)

# 7、初始化评估模块
# tru = Tru()
# tru.reset_database()

# 自动合并检索器